import numpy as np

class BicycleModelKalmanFilter:
    def __init__(self, dt, wheelbase=1, 
                 process_noise_std=[0.1, 0.1, 0.05, 0.02], 
                 measurement_noise_std=[0.1, 0.5, 0.1, 0.05]):
        """
        自行车模型卡尔曼滤波器
        
        参数:
        dt: 时间步长
        wheelbase: 轴距 (默认1米)
        process_noise_std: 过程噪声标准差数组 [x, y, v, psi]
        measurement_noise_std: 测量噪声标准差数组 [x, y, v, psi]
        """
        self.dt = dt
        self.L = wheelbase  # 轴距
        
        # 状态向量: [x, y, v, psi]^T
        # x: x位置, y: y位置, v: 速度, psi: 航向角
        self.state_dim = 4
        self.control_dim = 2  # [加速度, 转向角]
        
        # 初始状态和协方差矩阵
        self.x = np.zeros(self.state_dim)  # 状态估计
        self.P = np.eye(self.state_dim)    # 估计误差协方差
        
        # 验证噪声参数维度
        if len(process_noise_std) != self.state_dim:
            raise ValueError(f"过程噪声数组长度应为{self.state_dim}")
        if len(measurement_noise_std) != self.state_dim:
            raise ValueError(f"测量噪声数组长度应为{self.state_dim}")
        
        # 过程噪声协方差矩阵 - 基于不同状态变量的不确定性
        # [x位置噪声, y位置噪声, 速度噪声, 航向角噪声]
        self.Q = np.diag(np.array(process_noise_std)**2)
        
        # 测量噪声协方差矩阵 - 基于不同传感器的精度
        # [x位置传感器噪声, y位置传感器噪声, 速度传感器噪声, 航向角传感器噪声]
        self.R = np.diag(np.array(measurement_noise_std)**2)
        
        # 测量矩阵 (假设我们可以直接测量所有状态)
        self.H = np.eye(self.state_dim)
        
        # 存储历史数据用于分析
        self.history = {
            'states': [],
            'covariances': [],
            'controls': [],
            'measurements': [],
            'time': []
        }
    
    def reset(self, initial_state=None):
        """
        重置滤波器状态
        
        参数:
        initial_state: 初始状态向量 [x, y, v, psi]
        """
        if initial_state is None:
            self.x = np.zeros(self.state_dim)
        else:
            self.x = initial_state.copy()
        
        self.P = np.eye(self.state_dim)
        self.history = {
            'states': [],
            'covariances': [],
            'controls': [],
            'measurements': [],
            'time': []
        }
    
    def bicycle_model(self, state, control):
        """
        自行车模型的状态转移函数
        
        参数:
        state: 当前状态 [x, y, v, psi]
        control: 控制输入 [加速度, 转向角]
        
        返回:
        新状态
        """
        x, y, v, psi = state
        a, delta = control  # 加速度, 转向角
        
        # 防止速度为负时的不合理转向
        if abs(v) < 0.1:
            v_new = max(0, v + a * self.dt)
            psi_new = psi
        else:
            v_new = v + a * self.dt
            # 自行车运动学模型
            psi_new = psi + (v_new / self.L) * np.tan(delta) * self.dt # todo 这里的航向角计算有问题？应该改成和环境里的一样——减号？？？
        
        # 位置更新
        x_new = x + v_new * np.cos(psi_new) * self.dt
        y_new = y + v_new * np.sin(psi_new) * self.dt
        
        return np.array([x_new, y_new, v_new, psi_new])
    
    def jacobian_F(self, state, control):
        """
        计算状态转移函数的雅可比矩阵 (关于状态)
        """
        x, y, v, psi = state
        a, delta = control
        
        # 计算新速度
        v_new = v + a * self.dt
        
        F = np.eye(self.state_dim)
        
        if abs(v) < 0.1:
            # 低速情况下的简化雅可比矩阵
            F[0, 2] = self.dt * np.cos(psi)  # ∂x/∂v
            F[0, 3] = 0  # ∂x/∂psi
            F[1, 2] = self.dt * np.sin(psi)  # ∂y/∂v
            F[1, 3] = 0  # ∂y/∂psi
        else:
            # 完整雅可比矩阵
            beta = (v_new / self.L) * np.tan(delta)
            F[0, 2] = self.dt * np.cos(psi) - (v_new * self.dt**2 / self.L) * np.tan(delta) * np.sin(psi)
            F[0, 3] = -v_new * self.dt * np.sin(psi)
            F[1, 2] = self.dt * np.sin(psi) + (v_new * self.dt**2 / self.L) * np.tan(delta) * np.cos(psi)
            F[1, 3] = v_new * self.dt * np.cos(psi)
            F[3, 2] = (self.dt / self.L) * np.tan(delta)  # ∂psi/∂v
            F[3, 3] = 1  # ∂psi/∂psi
        
        return F
    
    def jacobian_G(self, state, control):
        """
        计算状态转移函数的雅可比矩阵 (关于控制)
        """
        x, y, v, psi = state
        a, delta = control
        
        v_new = v + a * self.dt
        
        G = np.zeros((self.state_dim, self.control_dim))
        
        if abs(v) < 0.1:
            # 低速情况
            G[2, 0] = self.dt  # ∂v/∂a
            G[0, 0] = 0.5 * self.dt**2 * np.cos(psi)  # ∂x/∂a
            G[1, 0] = 0.5 * self.dt**2 * np.sin(psi)  # ∂y/∂a
        else:
            # 完整情况
            G[0, 0] = 0.5 * self.dt**2 * np.cos(psi)  # ∂x/∂a
            G[0, 1] = -(v_new * self.dt**2 / (self.L * np.cos(delta)**2)) * np.sin(psi)  # ∂x/∂delta
            
            G[1, 0] = 0.5 * self.dt**2 * np.sin(psi)  # ∂y/∂a
            G[1, 1] = (v_new * self.dt**2 / (self.L * np.cos(delta)**2)) * np.cos(psi)  # ∂y/∂delta
            
            G[2, 0] = self.dt  # ∂v/∂a
            
            G[3, 1] = (v_new * self.dt) / (self.L * np.cos(delta)**2)  # ∂psi/∂delta
        
        return G
    
    def predict(self, control):
        """
        预测步骤
        
        参数:
        control: 控制输入 [加速度, 转向角]
        
        返回:
        预测后的状态
        """
        # 状态预测
        self.x = self.bicycle_model(self.x, control)
        
        # 计算雅可比矩阵
        F = self.jacobian_F(self.x, control)
        
        # 协方差预测
        self.P = F @ self.P @ F.T + self.Q
        
        # 记录历史
        self.history['states'].append(self.x.copy())
        self.history['covariances'].append(self.P.copy())
        self.history['controls'].append(control.copy())
        self.history['time'].append(len(self.history['time']) * self.dt)
        
        return self.x.copy()
    
    def update(self, measurement):
        """
        更新步骤
        
        参数:
        measurement: 观测向量 [x, y, v, psi]
        
        返回:
        更新后的状态和残差
        """
        # 计算卡尔曼增益
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)
        
        # 状态更新
        y = measurement - self.H @ self.x  # 残差
        self.x = self.x + K @ y
        
        # 协方差更新
        I = np.eye(self.state_dim)
        self.P = (I - K @ self.H) @ self.P
        
        # 记录测量值
        self.history['measurements'].append(measurement.copy())
        
        return self.x.copy(), y
    
    def filter_step(self, measurement, control):
        """
        完整的滤波步骤: 预测 + 更新
        
        参数:
        measurement: 观测向量 [x, y, v, psi]
        control: 控制输入 [加速度, 转向角]
        
        返回:
        估计的状态和残差
        """
        self.predict(control)
        estimated_state, residual = self.update(measurement)
        return estimated_state, residual
    
    def set_state(self, state):
        """设置状态"""
        self.x = state.copy()
    
    def get_state(self):
        """返回当前状态估计"""
        return self.x.copy()
    
    def get_covariance(self):
        """返回当前协方差矩阵"""
        return self.P.copy()
    
    # ignore——忽略，不做动态参数
    def set_noise_parameters(self, process_noise_std=None, measurement_noise_std=None):
        """
        动态设置噪声参数
        
        参数:
        process_noise_std: 新的过程噪声标准差数组
        measurement_noise_std: 新的测量噪声标准差数组
        """
        if process_noise_std is not None:
            if len(process_noise_std) != self.state_dim:
                raise ValueError(f"过程噪声数组长度应为{self.state_dim}")
            self.Q = np.diag(np.array(process_noise_std)**2)
        
        if measurement_noise_std is not None:
            if len(measurement_noise_std) != self.state_dim:
                raise ValueError(f"测量噪声数组长度应为{self.state_dim}")
            self.R = np.diag(np.array(measurement_noise_std)**2)


# 示例使用
if __name__ == "__main__":
    # 创建滤波器 - 使用数组形式的噪声参数
    kf = BicycleModelKalmanFilter(
        dt=0.1,
        wheelbase=2.5,
        process_noise_std=[0.1, 0.1, 0.05, 0.02],  # [x, y, v, psi]
        measurement_noise_std=[0.5, 0.5, 0.1, 0.05]  # [x, y, v, psi]
    )
    
    # 设置初始状态
    initial_state = np.array([0.0, 0.0, 10.0, 0.0])  # [x, y, v, psi]
    kf.set_state(initial_state)
    
    # 模拟一些控制输入和测量
    print("自行车模型卡尔曼滤波示例:")
    for i in range(10):
        # 控制输入: [加速度, 转向角]
        control = np.array([0.5, 0.1])  # 轻微加速和转向
        
        # 模拟测量 (带噪声)
        true_state = kf.get_state()
        measurement_noise = np.random.normal(0, [0.5, 0.5, 0.1, 0.05])
        measurement = true_state + measurement_noise
        
        # 执行滤波
        estimated_state, residual = kf.filter_step(measurement, control)
        
        print(f"步骤 {i+1}:")
        print(f"  控制: 加速度={control[0]:.2f}, 转向角={control[1]:.2f}")
        print(f"  测量: x={measurement[0]:.2f}, y={measurement[1]:.2f}, v={measurement[2]:.2f}, psi={measurement[3]:.2f}")
        print(f"  估计: x={estimated_state[0]:.2f}, y={estimated_state[1]:.2f}, v={estimated_state[2]:.2f}, psi={estimated_state[3]:.2f}")
        print(f"  残差: {residual}")
        print()